#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
FAISS检索器测试套件

测试所有FAISS-based检索器的功能、性能和稳定性
"""

import sys
import os
import time
import numpy as np
from typing import List, Dict, Any
from loguru import logger

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 在导入任何库之前强制设置环境变量
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_CO_URL_HOME"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.expanduser("~/.cache/huggingface")
os.environ["SENTENCE_TRANSFORMERS_HOME"] = os.path.expanduser("~/.cache/sentence_transformers")

# 强制导入配置，确保镜像环境变量已设置
try:
    import questionretrieval.config
    logger.info("🌐 使用 HuggingFace 国内镜像配置")
    logger.info(f"📁 缓存目录: {questionretrieval.config.HF_HUB_CACHE}")
    logger.info(f"🔗 镜像地址: {questionretrieval.config.HF_MIRROR_ENDPOINT}")
except ImportError:
    logger.warning("⚠️  无法导入配置模块，使用默认配置")

import pytest

# 测试数据
from questionretrieval.io.data_loader import load_questions

# 从questions.json加载真实数据
try:
    import os
    questions_path = os.path.join(os.path.dirname(__file__), '..', 'questionretrieval', 'io', 'questions.json')
    SAMPLE_QUESTIONS = load_questions(questions_path)
    if not SAMPLE_QUESTIONS:
        logger.warning("⚠️  无法从questions.json加载数据，使用默认测试数据")
        SAMPLE_QUESTIONS = [
            "什么是机器学习？",
            "如何学习Python编程？",
            "深度学习与机器学习的区别是什么？",
            "什么是神经网络？",
            "如何优化模型性能？"
        ]
except Exception as e:
    logger.warning(f"⚠️  加载questions.json失败: {e}，使用默认测试数据")
    SAMPLE_QUESTIONS = [
        "什么是机器学习？",
        "如何学习Python编程？",
        "深度学习与机器学习的区别是什么？",
        "什么是神经网络？",
        "如何优化模型性能？"
    ]

SAMPLE_QUERIES = [
    "大伙房水库的总库容是多少？",
    "大伙房水库的设计标准是什么？",
    "大伙房水库位于哪个流域？",
    "大伙房水库为哪些城市提供供水？",
    "大伙房水库的防洪调度原则是什么？"
]


def test_flat_index_retriever():
    """测试Flat Index检索器"""
    logger.info("\n=== 测试 FlatIndexRetriever ===")
    
    try:
        from questionretrieval.retrievers.faiss import FlatIndexRetriever
        
        # 测试初始化
        retriever = FlatIndexRetriever(
            model_name='all-MiniLM-L6-v2',
            metric='cosine',
            normalize_vectors=True
        )
        logger.info("✓ FlatIndexRetriever 初始化成功")
        
        # 测试加载候选问题
        retriever.load_candidates(SAMPLE_QUESTIONS)
        logger.info(f"✓ 成功加载 {len(SAMPLE_QUESTIONS)} 个候选问题")
        
        # 测试索引信息
        info = retriever.get_index_info()
        assert info['status'] == 'built'
        assert info['total_vectors'] == len(SAMPLE_QUESTIONS)
        assert info['metric'] == 'cosine'
        logger.info(f"✓ 索引信息: {info}")
        
        # 测试推荐功能
        results = retriever.recommend("什么是机器学习？", top_k=3)
        assert len(results) == 3
        assert all(isinstance(r, tuple) and len(r) == 3 for r in results)
        logger.info(f"✓ 推荐功能正常: {len(results)} 个结果")
        
        # 测试不同指标
        retriever_l2 = FlatIndexRetriever(metric='l2')
        retriever_l2.load_candidates(SAMPLE_QUESTIONS)
        results_l2 = retriever_l2.recommend("什么是机器学习？", top_k=3)
        assert len(results_l2) == 3
        logger.info("✓ L2距离指标测试成功")
        
        return True
        
    except ImportError as e:
        logger.warning(f"✗ FlatIndexRetriever 不可用: {e}")
        return False
    except Exception as e:
        logger.error(f"✗ FlatIndexRetriever 测试失败: {e}")
        return False


def test_ivf_index_retriever():
    """测试IVF Index检索器"""
    logger.info("\n=== 测试 IVFIndexRetriever ===")
    
    try:
        from questionretrieval.retrievers.faiss import IVFIndexRetriever
        
        # 测试初始化
        retriever = IVFIndexRetriever(
            model_name='all-MiniLM-L6-v2',
            nlist=5,  # 小数据集使用较小的nlist
            nprobe=2,
            metric='cosine'
        )
        logger.info("✓ IVFIndexRetriever 初始化成功")
        
        # 测试加载候选问题
        retriever.load_candidates(SAMPLE_QUESTIONS)
        logger.info(f"✓ 成功加载 {len(SAMPLE_QUESTIONS)} 个候选问题")
        
        # 测试索引信息
        info = retriever.get_index_info()
        assert info['status'] == 'built'
        assert info['total_vectors'] == len(SAMPLE_QUESTIONS)
        assert info['nlist'] == 5
        assert info['nprobe'] == 2
        logger.info(f"✓ 索引信息: {info}")
        
        # 测试推荐功能
        results = retriever.recommend("什么是机器学习？", top_k=3)
        assert len(results) == 3
        logger.info(f"✓ 推荐功能正常: {len(results)} 个结果")
        
        # 测试搜索参数调整
        retriever.set_search_params(nprobe=3)
        assert retriever.nprobe == 3
        logger.info("✓ 搜索参数调整成功")
        
        # 测试带nprobe参数的推荐
        results_custom = retriever.recommend("什么是机器学习？", top_k=3, nprobe=1)
        assert len(results_custom) == 3
        logger.info("✓ 自定义nprobe参数测试成功")
        
        return True
        
    except ImportError as e:
        logger.warning(f"✗ IVFIndexRetriever 不可用: {e}")
        return False
    except Exception as e:
        logger.error(f"✗ IVFIndexRetriever 测试失败: {e}")
        return False


def test_hnsw_index_retriever():
    """测试HNSW Index检索器"""
    logger.info("\n=== 测试 HNSWIndexRetriever ===")
    
    try:
        from questionretrieval.retrievers.faiss import HNSWIndexRetriever
        
        # 测试初始化
        retriever = HNSWIndexRetriever(
            model_name='all-MiniLM-L6-v2',
            M=8,  # 小数据集使用较小的M
            efConstruction=100,
            efSearch=20,
            metric='cosine'
        )
        logger.info("✓ HNSWIndexRetriever 初始化成功")
        
        # 测试加载候选问题
        retriever.load_candidates(SAMPLE_QUESTIONS)
        logger.info(f"✓ 成功加载 {len(SAMPLE_QUESTIONS)} 个候选问题")
        
        # 测试索引信息
        info = retriever.get_index_info()
        assert info['status'] == 'built'
        assert info['total_vectors'] == len(SAMPLE_QUESTIONS)
        assert info['M'] == 8
        assert info['efSearch'] == 20
        logger.info(f"✓ 索引信息: {info}")
        
        # 测试推荐功能
        results = retriever.recommend("什么是机器学习？", top_k=3)
        assert len(results) == 3
        logger.info(f"✓ 推荐功能正常: {len(results)} 个结果")
        
        # 测试搜索参数调整
        retriever.set_search_params(efSearch=30)
        assert retriever.efSearch == 30
        logger.info("✓ 搜索参数调整成功")
        
        # 测试带efSearch参数的推荐
        results_custom = retriever.recommend("什么是机器学习？", top_k=3, efSearch=10)
        assert len(results_custom) == 3
        logger.info("✓ 自定义efSearch参数测试成功")
        
        return True
        
    except ImportError as e:
        logger.warning(f"✗ HNSWIndexRetriever 不可用: {e}")
        return False
    except Exception as e:
        logger.error(f"✗ HNSWIndexRetriever 测试失败: {e}")
        return False


def test_faiss_index_persistence():
    """测试FAISS索引的保存和加载功能"""
    logger.info("\n=== 测试 FAISS索引持久化 ===")
    
    try:
        from questionretrieval.retrievers.faiss import FlatIndexRetriever
        import tempfile
        import os
        
        # 创建临时文件
        with tempfile.NamedTemporaryFile(suffix='.faiss', delete=False) as tmp_file:
            tmp_path = tmp_file.name
        
        try:
            # 创建并保存索引
            retriever = FlatIndexRetriever()
            retriever.load_candidates(SAMPLE_QUESTIONS)
            retriever.save_index(tmp_path)
            logger.info("✓ 索引保存成功")
            
            # 创建新检索器并加载索引
            new_retriever = FlatIndexRetriever()
            new_retriever.load_index(tmp_path)
            new_retriever.questions = SAMPLE_QUESTIONS  # 需要手动设置questions
            new_retriever.scores = [1.0] * len(SAMPLE_QUESTIONS)
            
            # 测试加载后的功能
            results = new_retriever.recommend("什么是机器学习？", top_k=3)
            assert len(results) == 3
            logger.info("✓ 索引加载和功能测试成功")
            
        finally:
            # 清理临时文件
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)
        
        return True
        
    except Exception as e:
        logger.error(f"✗ 索引持久化测试失败: {e}")
        return False


def test_faiss_performance_benchmark():
    """测试FAISS检索器性能基准"""
    logger.info("\n=== 测试 FAISS性能基准 ===")
    
    try:
        from questionretrieval.retrievers.faiss import FlatIndexRetriever, IVFIndexRetriever, HNSWIndexRetriever
        
        # 创建更大的测试数据集
        large_questions = SAMPLE_QUESTIONS * 5  # 50个问题
        
        retrievers = {
            'Flat': FlatIndexRetriever(),
            'IVF': IVFIndexRetriever(nlist=10, nprobe=3),
            'HNSW': HNSWIndexRetriever(M=8, efConstruction=100, efSearch=20)
        }
        
        results = {}
        
        for name, retriever in retrievers.items():
            try:
                start_time = time.time()
                retriever.load_candidates(large_questions)
                build_time = time.time() - start_time
                
                # 测试查询性能
                start_time = time.time()
                for query in SAMPLE_QUERIES[:2]:  # 使用2个查询进行测试
                    retriever.recommend(query, top_k=5)
                query_time = time.time() - start_time
                
                results[name] = {
                    'build_time': build_time,
                    'query_time': query_time,
                    'queries_per_second': 2 / query_time if query_time > 0 else 0
                }
                
                logger.info(f"✓ {name} 性能测试完成: 构建{build_time:.3f}s, 查询{query_time:.3f}s")
                
            except Exception as e:
                logger.error(f"✗ {name} 性能测试失败: {e}")
                results[name] = {'error': str(e)}
        
        # 验证至少有一个检索器成功
        successful = [r for r in results.values() if 'error' not in r]
        if successful:
            logger.info("✓ 性能基准测试完成")
            return True
        else:
            logger.error("✗ 所有检索器性能测试失败")
            return False
            
    except Exception as e:
        logger.error(f"✗ 性能基准测试失败: {e}")
        return False


def test_all_faiss_retrievers():
    """测试所有FAISS检索器"""
    logger.info("\n" + "="*60)
    logger.info("FAISS检索器测试套件")
    logger.info("="*60)
    
    # 定义测试函数
    tests = [
        ("FlatIndexRetriever", test_flat_index_retriever),
        ("IVFIndexRetriever", test_ivf_index_retriever),
        ("HNSWIndexRetriever", test_hnsw_index_retriever),
        ("IndexPersistence", test_faiss_index_persistence),
        ("PerformanceBenchmark", test_faiss_performance_benchmark)
    ]
    
    results = {}
    
    for test_name, test_func in tests:
        try:
            results[test_name] = test_func()
        except Exception as e:
            logger.error(f"✗ {test_name} 测试异常: {e}")
            results[test_name] = False
    
    # 输出结果总结
    logger.info("\n" + "-"*50)
    logger.info("FAISS测试总结")
    logger.info("-"*50)
    
    success_count = 0
    for name, success in results.items():
        status = "✓ 通过" if success else "✗ 失败"
        logger.info(f"{name:25} : {status}")
        if success:
            success_count += 1
    
    logger.info(f"\n总计: {success_count}/{len(tests)} 个测试通过")
    
    if success_count > 0:
        logger.info("🎉 恭喜! 部分FAISS检索器测试通过")
        logger.info("💡 通过的检索器可以在项目中使用")
    else:
        logger.warning("⚠️  所有FAISS检索器测试失败")
        logger.info("请检查依赖: pip install faiss-cpu sentence-transformers")
    
    return results


if __name__ == "__main__":
    test_all_faiss_retrievers()